#!/usr/bin/env python3
"""
简单测试脚本 - 两步投资分析流程
"""

import json
import os
import sys
from pathlib import Path
from datetime import datetime

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from app.services.core_service import CoreService

def test_investment_analysis():
    """测试两步投资分析流程 - 保存详细过程结果"""
    # 设置环境变量
    os.environ["OPENAI_API_KEY"] = "sk-gX6NNINB54pDBGRka36jg33fR1fnSOwfu21uC3Wnjmiwv3KB"
    os.environ["OPENAI_API_BASE"] = "https://api.chatanywhere.tech/v1"
    
    # 加载新闻数据
    with open('data/processed_news.json', 'r', encoding='utf-8') as f:
        all_news = json.load(f)
    
    # 初始化核心服务
    core_service = CoreService()
    core_service.initialize()
    
    # 创建统一结果文件
    os.makedirs('result', exist_ok=True)
    result_file = 'result/detailed_analysis_results.json'
    
    # 如果文件存在，读取现有结果
    all_results = []
    if os.path.exists(result_file):
        with open(result_file, 'r', encoding='utf-8') as f:
            all_results = json.load(f)
    
    # 测试每条新闻
    for i, news_data in enumerate(all_news, 1):
        # 分析新闻
        result = core_service.analyze_news(news_data)
        
        # 创建详细的过程结果
        detailed_result = {
            "batch_index": i,
            "batch_timestamp": str(datetime.now()),
            "news_info": {
                "id": news_data.get('news_id', ''),
                "title": news_data.get('title', ''),
                "content": news_data.get('content', '')[:200] + "..." if len(news_data.get('content', '')) > 200 else news_data.get('content', '')
            },
            "analysis_process": {
                "step1_token_extraction": result.get('token_extraction', {}),
                "step2_entities": result.get('entities', {}),
                "step3_queries": result.get('queries', []),
                "step4_knowledge_retrieval": result.get('knowledge', []),
                "step5_price_analysis": result.get('price_analysis', {}),
                "final_status": result.get('status', '')
            },
            "rag_details": {
                "retrieved_knowledge_count": len(result.get('knowledge', [])),
                "knowledge_sources": [k[:100] + "..." if len(k) > 100 else k for k in result.get('knowledge', [])],
                "query_generation": [q.get('query', '') for q in result.get('queries', [])],
                "entity_extraction": result.get('entities', {})
            },
            "learning_context": {
                "enhanced_prompt": "基于学习到的成功模式，特别关注关键字",
                "confidence_adjustment": "根据历史成功经验调整置信度"
            }
        }
        
        # 追加到结果列表
        all_results.append(detailed_result)
    
    # 保存所有结果到统一文件
    with open(result_file, 'w', encoding='utf-8') as f:
        json.dump(all_results, f, ensure_ascii=False, indent=2)

def test_learning():
    """测试学习功能"""
    # 设置环境变量
    os.environ["OPENAI_API_KEY"] = "sk-gX6NNINB54pDBGRka36jg33fR1fnSOwfu21uC3Wnjmiwv3KB"
    os.environ["OPENAI_API_BASE"] = "https://api.chatanywhere.tech/v1"
    
    # 加载新闻数据
    with open('data/processed_news.json', 'r', encoding='utf-8') as f:
        all_news = json.load(f)
    
    # 加载实际结果
    with open('data/actual_results.json', 'r', encoding='utf-8') as f:
        actual_results = json.load(f)
    
    # 初始化核心服务
    core_service = CoreService()
    core_service.initialize()
    
    # 测试学习
    for news_data, actual_result in zip(all_news, actual_results):
        # 分析新闻
        result = core_service.analyze_news(news_data)
        
        if result['status'] == 'success' and 'price_analysis' in result:
            # 模拟学习
            core_service.learn_from_result(
                news_data, 
                result['price_analysis'], 
                actual_result['actual_result']
            )

if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "learn":
        test_learning()
    else:
        test_investment_analysis()
